{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最小MNIST神经网络模型例子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.\n",
      "It can be downloaded at https://aka.ms/vs/16/release/vc_redist.x64.exe\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 假设我们有两个模型类 Model1 和 Model2\n",
    "class Model1(nn.Module):\n",
    "    # 定义Model1的架构\n",
    "    pass\n",
    "\n",
    "class Model2(nn.Module):\n",
    "    # 定义Model2的架构\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 通用测试函数\n",
    "def evaluate_model(\n",
    "    model, data_loader, criterion, device=\"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "):\n",
    "    model.eval()  # 设置为评估模式\n",
    "    total_correct = 0\n",
    "    total_samples = 0\n",
    "    total_loss = 0.0\n",
    "\n",
    "    with torch.no_grad():  # 不需要计算梯度\n",
    "        for images, labels in data_loader:\n",
    "            images, labels = images.to(device), labels.to(\n",
    "                device\n",
    "            )  # 将数据移动到指定设备\n",
    "            outputs = model(images)  # 获取模型预测\n",
    "            loss = criterion(outputs, labels)  # 计算损失\n",
    "            total_loss += loss.item() * images.size(0)  # 累加损失\n",
    "            _, predicted = torch.max(outputs, 1)  # 获取预测类别\n",
    "            total_correct += (predicted == labels).sum().item()  # 累加正确预测的数量\n",
    "            total_samples += labels.size(0)  # 累加样本总数\n",
    "\n",
    "    # 计算平均损失和准确率\n",
    "    avg_loss = total_loss / total_samples\n",
    "    accuracy = total_correct / total_samples\n",
    "\n",
    "    return avg_loss, accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据预处理和数据加载器\n",
    "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])\n",
    "testset = datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=False, transform=transform)\n",
    "testloader = DataLoader(testset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 实例化模型和损失函数\n",
    "model1 = Model1().to('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model2 = Model2().to('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 评估模型\n",
    "avg_loss1, accuracy1 = evaluate_model(model1, testloader, criterion)\n",
    "avg_loss2, accuracy2 = evaluate_model(model2, testloader, criterion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 打印结果\n",
    "print(f'Model 1: Average Loss = {avg_loss1:.4f}, Accuracy = {accuracy1:.2f}')\n",
    "print(f'Model 2: Average Loss = {avg_loss2:.4f}, Accuracy = {accuracy2:.2f}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
